#include <algorithm>
#include "position_model.h"

PositionModel::PositionModel(int L, int n_obs, int max_length)
{
  motif_length_ = L;
  n_observation_ = n_obs;
  max_seq_length_ = max_length;
  rand(var_);
}

void PositionModel::rand(variable_t& var)
{
  var.eprob.clear();
  for (int i = 0; i <= motif_length_; i++) {
    var.eprob.push_back(arma::randu<arma::vec>(n_observation_));
    var.eprob.back() /= arma::sum(var.eprob.back());  // sum to unit
  }

  var.pprob = arma::randu<arma::vec>(max_seq_length_);
  var.pprob /= arma::sum(var.pprob);

  var.z = var.ACC_pprob = var.pprob;
  var.ACC_eprob = var.eprob;
}

void PositionModel::print_params() const
{
  arma::mat eprob(n_observation_, motif_length_+1);
  for (int i = 0; i <= motif_length_; i++) {
    eprob.col(i) = var_.eprob[i];
  }
  std::cout << "eprob = " << std::endl << eprob << "pprob = ";
  for (int i = 0; i < max_seq_length_; i++)
    std::cout << "p(" << i << ")=" << var_.pprob[i] << " ";
  std::cout << std::endl;
}

void PositionModel::train(const vseq_t& data, int max_iter)
{
  double max_s = -std::numeric_limits<double>::infinity();
  variable_t max_var;

  for (int n_rand = 0; n_rand < 10; n_rand++) {
    std::cout << "==================== n_rand = " << n_rand << " ========================= " << std::endl;
    double old_s = -1e6;
    for (int iter = 0; iter < max_iter; iter++) {
      std::cout << "--------- iter = " << iter << " -------------- ";
      //print_params();

      start_accumulate(var_);
      double s = 0;
      for (unsigned int seq_idx = 0; seq_idx < data.size(); seq_idx++) {
	const sequence_t& seq = data[seq_idx];
	compute_z(seq, var_);
	accumulate(seq, var_);
	s += var_.seq_prob;
      }
      //std::cout << " ---- " << var_.ACC_pprob[45] << " ---- " << var_.z[45]  <<  " z=" << arma::trans(var_.z) << std::endl;
      std::cout << "likelihood = " << s << std::endl;
      if (s > max_s) {
	max_s = s;
	max_var = var_;
      }
      reestimate(var_);

      if ( fabs(s-old_s) / fabs(old_s) < 1e-6 ) {
	std::cout << "Changes smaller than tolerance. Terminate" << std::endl;
	break;
      }
      old_s = s;
    }
    rand(var_);
  }
  var_ = max_var;
}

void PositionModel::start_accumulate(variable_t& var)
{
  var.ACC_pprob.fill(log(0));
  var.z.fill(log(0));
  for (int i = 0; i <= motif_length_; i++) var.ACC_eprob[i].fill(log(0));
}

void PositionModel::compute_z(const sequence_t& seq, variable_t& var)
{
  for (int i = 0; i < max_seq_length_ && i < (int) seq.size(); i++) {
    if (i > (int) seq.size() - motif_length_) {
      var.z[i] = log(0);
      continue;
    }
    
    // i <= max_seq_length_ - motif_length_
    var.z[i] = log(var.pprob[i]);
    for (int pos = 0; pos < max_seq_length_ && pos < (int) seq.size(); pos++) {
      if (pos >= i && pos < i+motif_length_)
	var.z[i] += log(var.eprob[pos-i+1][seq[pos]]); // motif emission at position pos-i+1 of value seq[pos]
      else
	var.z[i] += log(var.eprob[0][seq[pos]]); // background emission of value seq[pos]
    }
  }
  var.seq_prob = logsum_vec(var.z);
  for (int i = 0; i < max_seq_length_ && i < (int) seq.size(); i++) 
    var.z[i] -= var.seq_prob;
}

void PositionModel::accumulate(const sequence_t& seq, variable_t& var)
{
  for (int i = 0; i < max_seq_length_ && i < (int) seq.size(); i++) {
    if (i > (int) seq.size() - motif_length_) {
      continue;
    }
    double &s = var.ACC_pprob[i];
    s = logadd(s, var.z[i]);
    for (int pos = 0; pos < max_seq_length_ && pos < (int) seq.size(); pos++) {
      if (pos >= i && pos < i+motif_length_) {
	double &s = var.ACC_eprob[pos-i+1][seq[pos]];
	s = logadd(s, var.z[i]);
      }
      else {
	double &s = var.ACC_eprob[0][seq[pos]];
	s = logadd(s, var.z[i]);
      }
    }
  }
}

void PositionModel::reestimate(variable_t& var)
{
  double s = logsum_vec(var.ACC_pprob);
  for (int i = 0; i < max_seq_length_; i++) {
    if (s == -std::numeric_limits<double>::infinity())
      var.pprob[i] = 0;
    else
      var.pprob[i] = exp(var.ACC_pprob[i] - s);
  }

  for (int i = 0; i <= motif_length_; i++ ) {
    double s = logsum_vec(var.ACC_eprob[i]);
    for (int j = 0; j < n_observation_; j++) {
      if (s == -std::numeric_limits<double>::infinity())
	var.eprob[i][j] = 0;
      else
	var.eprob[i][j] = exp(var.ACC_eprob[i][j] - s);
    }
  }
}

int PositionModel::map_pos(const sequence_t& seq)
{
  var_.z.fill(log(0));
  compute_z(seq, var_);
  int p = 0;
  for (int i = 0; i < max_seq_length_; i++)
    if (var_.z[i] > var_.z[p]) p = i;
  return p;
}
